import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as patches

class Component:
    def __init__(self, name, width, height, connections=None):
        self.name = name
        self.width = width
        self.height = height
        self.connections = connections if connections else []
        self.x = np.random.rand() * 100  # 初始随机位置
        self.y = np.random.rand() * 100
        self.locked = False #元器件是否锁定坐标

    def add_connection(self, component):
        self.connections.append(component)

class PCBPlacer:
    def __init__(self, width=100, height=100):
        self.pcb_width = width
        self.pcb_height = height
        self.components = []
        self.attraction_coeff = 0.01    # 吸引力系数
        self.repulsion_coeff = 10000    # 斥力系数
        self.edge_repulsion = 100       # 边缘斥力系数
        self.damping = 0.85             # 阻尼系数
        self.max_speed = 2.0            # 最大移动速度
        
    def add_component(self, component):
        self.components.append(component)
        
    def calculate_forces(self):
        forces = np.zeros((len(self.components), 2))

        
        # 计算连接之间的吸引力
        for i, comp in enumerate(self.components):
            for connected in comp.connections:
                j = self.components.index(connected)
                dx = self.components[j].x - comp.x
                dy = self.components[j].y - comp.y
                distance = np.hypot(dx, dy)  #计算dx，dy为直角边的三角形的斜边长度
                
                if distance > 0:
                    force = self.attraction_coeff * distance
                    forces[i][0] += force * dx / distance
                    forces[i][1] += force * dy / distance
                    
        # 计算所有元件之间的斥力
        for i in range(len(self.components)):
            for j in range(i+1, len(self.components)):
                comp1 = self.components[i]
                comp2 = self.components[j]
                
                dx = comp2.x - comp1.x
                dy = comp2.y - comp1.y
                distance = np.hypot(dx, dy)
                
                min_distance = (comp1.width + comp2.width)/2 + 5
                if distance < min_distance:
                    force = self.repulsion_coeff / max(distance, 0.1)**2
                    forces[i][0] -= force * dx / distance
                    forces[i][1] -= force * dy / distance
                    forces[j][0] += force * dx / distance
                    forces[j][1] += force * dy / distance
                    
        # 计算边缘斥力
        for i, comp in enumerate(self.components):
            # 左边界
            force = self.edge_repulsion / (comp.x + 1)
            forces[i][0] += force
            # 右边界
            force = self.edge_repulsion / (self.pcb_width - comp.x + 1)
            forces[i][0] -= force
            # 下边界
            force = self.edge_repulsion / (comp.y + 1)
            forces[i][1] += force
            # 上边界
            force = self.edge_repulsion / (self.pcb_height - comp.y + 1)
            forces[i][1] -= force
            
        return forces
    
    def update_positions(self, forces):
        for i, comp in enumerate(self.components):
            if comp.locked == True:
                continue
            # 限制最大速度
            fx = max(min(forces[i][0], self.max_speed), -self.max_speed)
            fy = max(min(forces[i][1], self.max_speed), -self.max_speed)
            
            comp.x += fx * self.damping
            comp.y += fy * self.damping
            
            # 保持元件在PCB边界内
            comp.x = max(comp.width/2, min(self.pcb_width - comp.width/2, comp.x))
            comp.y = max(comp.height/2, min(self.pcb_height - comp.height/2, comp.y))
    
    def visualize(self):
        fig, ax = plt.subplots(figsize=(10, 10))
        ax.set_xlim(0, self.pcb_width)
        ax.set_ylim(0, self.pcb_height)
        
        # 绘制PCB边界
        border = patches.Rectangle((0,0), self.pcb_width, self.pcb_height, 
                                 linewidth=2, edgecolor='black', facecolor='none')
        ax.add_patch(border)
        
        # 绘制元件
        for comp in self.components:
            rect = patches.Rectangle((comp.x-comp.width/2, comp.y-comp.height/2),
                                   comp.width, comp.height,
                                   linewidth=1, edgecolor='blue', facecolor='lightblue')
            ax.add_patch(rect)
            plt.text(comp.x, comp.y, comp.name, ha='center', va='center')
            
            # 绘制连接线
            for connected in comp.connections:
                end = connected
                plt.plot([comp.x, end.x], [comp.y, end.y], 'gray', linewidth=0.5)
                
        plt.gca().set_aspect('equal', adjustable='box')
        plt.show()

# 示例用法
if __name__ == "__main__":


    arr = []

    for x in range(10):
        arr.append(Component(str(x), np.random.randint(low=3, high=10), np.random.randint(low=3, high=10)))

    
    # 创建元件
    comp1 = Component("U1", np.random.randint(low=3, high=10), np.random.randint(low=3, high=10))
    comp2 = Component("U2", np.random.randint(low=3, high=10), np.random.randint(low=3, high=10))
    comp3 = Component("U3", np.random.randint(low=3, high=10), np.random.randint(low=3, high=10))
    comp4 = Component("U4", np.random.randint(low=3, high=10), np.random.randint(low=3, high=10))
    comp5 = Component("U5", np.random.randint(low=3, high=10), np.random.randint(low=3, high=10))
    comp6 = Component("U6", np.random.randint(low=3, high=10), np.random.randint(low=3, high=10))
    comp7 = Component("U7", np.random.randint(low=3, high=10), np.random.randint(low=3, high=10))
    comp8 = Component("U8", np.random.randint(low=3, high=10), np.random.randint(low=3, high=10))
    comp1.locked = True
    # comp8.locked = True

    # 添加连接关系
    comp1.add_connection(comp2)
    comp1.add_connection(comp3)
    comp2.add_connection(comp3)
    comp2.add_connection(comp4)
    comp3.add_connection(comp4)
    comp4.add_connection(comp2)
    comp5.add_connection(comp1)
    comp6.add_connection(comp7)
    comp7.add_connection(comp5)
    comp8.add_connection(comp4)
    
    # 创建布局器
    placer = PCBPlacer(200, 200)
    placer.add_component(comp1)
    placer.add_component(comp2)
    placer.add_component(comp3)
    placer.add_component(comp4)
    placer.add_component(comp5)
    placer.add_component(comp6)
    placer.add_component(comp7)
    placer.add_component(comp8)
    
    placer.visualize()
    # 迭代布局
    for n in range(3000):
        forces = placer.calculate_forces()
        placer.update_positions(forces)
        if (n%100) == 0:
            placer.visualize()

    
    # 可视化结果
    placer.visualize()